/* SPDX-License-Identifier: GPL-2.0 */
/*
 * Copyright (c) 2024
 */
#include <stdio.h>
#include <unistd.h>
#include <signal.h>
#include <string.h>
#include <libgen.h>
#include <bpf/bpf.h>
#include <scx/common.h>
#include <stdlib.h>
#include <pthread.h>
#include <time.h>
#include <dirent.h>
#include <ctype.h>
#include <sys/types.h>
#include <sstream>
#include <fstream>
#include <vector>
#include <algorithm>
#include <string>
#include <iomanip>
#include <unordered_map>
#include <unordered_set>
#include "pcerrc.h"
#include "pmu.h"
#include "symbol.h"
#include "scx_hello.h"
#include "scx_hello.bpf.skel.h"

const char help_fmt[] =
"一个演示用户态和内核态BPF通信的调度器，支持传递LLC缓存未命中率信息。\n"
"\n"
"用法: %s [-f] [-v] [-m 消息] [-i 采集间隔(秒)] [-t 采集次数]\n"
"\n"
"  -f            使用FIFO调度而不是加权vtime调度\n"
"  -v            打印libbpf调试消息\n"
"  -m 消息       发送消息到内核态BPF程序\n"
"  -i 间隔       LLC缓存未命中率采集间隔，默认为5秒\n"
"  -t 次数       LLC缓存未命中率采集次数，默认为0（无限循环）\n"
"  -h            显示帮助信息并退出\n";

static bool verbose;
static volatile int exit_req;
static char user_message[256] = "Hello, Kernel!";
static int llc_collect_interval = 5; // 默认5秒采集一次
static int llc_collect_times = 0;    // 默认无限循环采集
static pthread_t llc_thread;
static struct scx_hello *g_skel = NULL;

// 线程信息结构
typedef struct {
    int tid;
    unsigned llc_miss;
    unsigned llc_access;
    double miss_ratio;
} thread_llc_info_t;

// 线程信息比较函数，用于排序
static int compare_thread_llc(const void *a, const void *b) {
    thread_llc_info_t *t1 = (thread_llc_info_t *)a;
    thread_llc_info_t *t2 = (thread_llc_info_t *)b;
    
    // 按未命中率降序排序
    if (t1->miss_ratio > t2->miss_ratio) return -1;
    if (t1->miss_ratio < t2->miss_ratio) return 1;
    return 0;
}

// 检查字符串是否为数字
static int is_numeric(const char *str) {
    while (*str) {
        if (!isdigit(*str))
            return 0;
        str++;
    }
    return 1;
}

// 获取进程名称
static void get_proc_name(int pid, char *name, size_t size) {
    char path[256];
    FILE *fp;
    
    snprintf(path, sizeof(path), "/proc/%d/comm", pid);
    fp = fopen(path, "r");
    if (fp) {
        if (fgets(name, size, fp)) {
            // 移除换行符
            size_t len = strlen(name);
            if (len > 0 && name[len-1] == '\n') {
                name[len-1] = '\0';
            }
        } else {
            snprintf(name, size, "unknown");
        }
        fclose(fp);
    } else {
        snprintf(name, size, "unknown");
    }
}

// 获取进程的线程列表
static int get_process_threads(int pid, int **threads) {
    char task_path[256];
    DIR *dir;
    struct dirent *entry;
    int count = 0;
    int *tids = NULL;
    int capacity = 10;
    
    snprintf(task_path, sizeof(task_path), "/proc/%d/task", pid);
    
    dir = opendir(task_path);
    if (!dir) {
        return 0;
    }
    
    // 分配初始内存
    tids = (int *)malloc(capacity * sizeof(int));
    if (!tids) {
        closedir(dir);
        return 0;
    }
    
    while ((entry = readdir(dir)) != NULL) {
        if (is_numeric(entry->d_name)) {
            int tid = atoi(entry->d_name);
            
            // 如果数组已满，扩容
            if (count >= capacity) {
                capacity *= 2;
                int *new_tids = (int *)realloc(tids, capacity * sizeof(int));
                if (!new_tids) {
                    free(tids);
                    closedir(dir);
                    return 0;
                }
                tids = new_tids;
            }
            
            tids[count++] = tid;
        }
    }
    
    closedir(dir);
    *threads = tids;
    return count;
}

// 解析CPU范围字符串，如"0-3,5"
static std::vector<int> parse_cpu_range(const std::string &range_str) {
    std::vector<int> cpus;
    std::stringstream ss(range_str);
    std::string part;

    while(std::getline(ss, part, ',')) {
        size_t hyphen_pos = part.find("-");
        if (hyphen_pos != std::string::npos) {
            int start = std::stoi(part.substr(0, hyphen_pos));
            int end = std::stoi(part.substr(hyphen_pos + 1));
            if (start > end) {
                fprintf(stderr, "Invalid CPU range: %s\n", part.c_str());
                continue;
            }
            for (int i = start; i <= end; ++i) {
                cpus.push_back(i);
            }
        } else {
            cpus.push_back(std::stoi(part));
        }
    }

    std::sort(cpus.begin(), cpus.end());
    cpus.erase(std::unique(cpus.begin(), cpus.end()), cpus.end());
    return cpus;
}

// 获取线程的CPU亲和性列表
static std::string get_cpu_affinity_list(int tid) {
    std::string path = "/proc/" + std::to_string(tid) + "/status";
    std::ifstream in(path);
    if (!in.is_open()) {
        fprintf(stderr, "Not found: %s\n", path.c_str());
        return "";
    }
    std::string line;
    const std::string target_key = "Cpus_allowed_list:";
    while (std::getline(in, line)) {
        if (line.find(target_key) == 0) {
            size_t pos = line.find("\t");
            if (pos == std::string::npos)
                pos = target_key.length();
            return line.substr(pos + 1);
        }
    }
    return "";
}

// 使用libkperf收集线程的LLC缓存未命中率
static int collect_thread_llc_miss_ratio(thread_llc_info_t *thread_info, int max_threads) {
    DIR *proc_dir;
    struct dirent *entry;
    int thread_count = 0;
    std::unordered_map<int, std::pair<unsigned, unsigned>> tid_data; // tid -> (llc_miss, llc_access)

    // 打开/proc目录
    proc_dir = opendir("/proc");
    if (!proc_dir) {
        perror("无法打开/proc目录");
        return 0;
    }
    
    // 遍历所有进程
    while ((entry = readdir(proc_dir)) != NULL && thread_count < max_threads) {
        // 只处理数字名称的目录（进程ID）
        if (is_numeric(entry->d_name)) {
            int pid = atoi(entry->d_name);
            int *threads = NULL;
            int thread_num = 0;
            
            // 获取进程的所有线程
            thread_num = get_process_threads(pid, &threads);
            
            // 处理每个线程
            for (int i = 0; i < thread_num && thread_count < max_threads; i++) {
                int tid = threads[i];
                std::string range_str = get_cpu_affinity_list(tid);
                if (range_str.empty()) {
                    continue;
                }
                
                std::vector<int> cpu_list = parse_cpu_range(range_str);
                if (cpu_list.empty()) {
                    continue;
                }
                
                // 使用libkperf采集LLC缓存未命中率
                char* evt_list[2];
                evt_list[0] = (char*)"r33"; // LLC缓存未命中事件
                evt_list[1] = (char*)"r32"; // LLC缓存访问事件
                
                PmuAttr attr = {0};
                attr.evtList = evt_list;
                attr.numEvt = 2;
                attr.pidList = &tid;
                attr.numPid = 1;
                attr.cpuList = cpu_list.data();
                attr.numCpu = cpu_list.size();
                
                int pd = PmuOpen(COUNTING, &attr);
                if (pd == -1) {
                    fprintf(stderr, "PmuOpen failed for tid %d\n", tid);
                    fprintf(stderr, "error msg: %s\n", Perror());
                    continue;
                }
                
                PmuEnable(pd);
                usleep(100000); // 采样100ms
                
                PmuData* pmu_data = nullptr;
                int len = PmuRead(pd, &pmu_data);
                if (len == -1) {
                    fprintf(stderr, "PmuRead failed: %s\n", Perror());
                    PmuDisable(pd);
                    PmuClose(pd);
                    continue;
                }
                
                // 处理采集到的数据
                for (int j = 0; j < len; ++j) {
                    PmuData& data = pmu_data[j];
                    if (strcmp(data.evt, "r33") == 0) {
                        tid_data[tid].first += data.count;
                    }
                    if (strcmp(data.evt, "r32") == 0) {
                        tid_data[tid].second += data.count;
                    }
                }
                
                PmuDataFree(pmu_data);
                PmuDisable(pd);
                PmuClose(pd);
            }
            
            if (threads) {
                free(threads);
            }
        }
    }
    
    closedir(proc_dir);
    
    // 将采集到的数据转换为thread_llc_info_t结构
    for (const auto& entry : tid_data) {
        if (thread_count >= max_threads) {
            break;
        }
        
        thread_info[thread_count].tid = entry.first;
        thread_info[thread_count].llc_miss = entry.second.first;
        thread_info[thread_count].llc_access = entry.second.second;
        thread_info[thread_count].miss_ratio = 
            (entry.second.second > 0) ? 
            ((double)entry.second.first / entry.second.second * 100.0) : 0.0;
        
        thread_count++;
    }
    
    // 按LLC缓存未命中率排序
    qsort(thread_info, thread_count, sizeof(thread_llc_info_t), compare_thread_llc);
    
    return thread_count;
}

// 发送LLC缓存未命中率信息到内核
static void send_llc_miss_info_to_kernel(struct scx_hello *skel, thread_llc_info_t *thread_info, int count) {
    u32 key = 0;
    struct scx_hello_llc_info llc_info;
    int llc_map_fd;
    int i;
    
    llc_map_fd = bpf_map__fd(skel->maps.llc_miss_info);
    if (llc_map_fd < 0) {
        fprintf(stderr, "无法获取llc_miss_info映射的文件描述符\n");
        return;
    }
    
    // 先读取当前信息
    if (bpf_map_lookup_elem(llc_map_fd, &key, &llc_info) < 0) {
        fprintf(stderr, "无法读取LLC缓存未命中率信息映射\n");
        return;
    }
    
    // 更新信息
    llc_info.count = (count > MAX_TOP_THREADS) ? MAX_TOP_THREADS : count;
    for (i = 0; i < llc_info.count; i++) {
        llc_info.threads[i].tid = thread_info[i].tid;
        llc_info.threads[i].miss_ratio = thread_info[i].miss_ratio;
    }
    llc_info.new_data = true;
    
    // 写回映射
    if (bpf_map_update_elem(llc_map_fd, &key, &llc_info, 0) < 0) {
        fprintf(stderr, "无法更新LLC缓存未命中率信息映射\n");
        return;
    }
    
    printf("已发送%d个线程的LLC缓存未命中率信息到内核\n", llc_info.count);
    for (i = 0; i < llc_info.count; i++) {
        char proc_name[64] = {0};
        get_proc_name(llc_info.threads[i].tid, proc_name, sizeof(proc_name));
        printf("  TID: %d (%s), Miss Ratio: %.2f%%, Miss: %u, Access: %u\n", 
               thread_info[i].tid, 
               proc_name,
               thread_info[i].miss_ratio,
               thread_info[i].llc_miss,
               thread_info[i].llc_access);
    }
}

// LLC缓存未命中率采集线程
static void *llc_collect_thread(void *arg) {
    thread_llc_info_t thread_info[1000]; // 假设最多收集1000个线程
    int count;
    int collect_count = 0;
    
    while (!exit_req && (llc_collect_times == 0 || collect_count < llc_collect_times)) {
        // 收集线程的LLC缓存未命中率
        count = collect_thread_llc_miss_ratio(thread_info, 1000);
        
        // 发送top10线程信息到内核
        if (count > 0 && g_skel) {
            send_llc_miss_info_to_kernel(g_skel, thread_info, count);
        }
        
        collect_count++;
        
        // 等待指定的间隔时间
        sleep(llc_collect_interval);
    }
    
    return NULL;
}

static int libbpf_print_fn(enum libbpf_print_level level, const char *format, va_list args)
{
    if (level == LIBBPF_DEBUG && !verbose)
        return 0;
    return vfprintf(stderr, format, args);
}

static void sigint_handler(int sig)
{
    exit_req = 1;
}

static void read_stats(struct scx_hello *skel, __u64 *stats)
{
    int nr_cpus = libbpf_num_possible_cpus();
    __u64 cnts[2][nr_cpus];
    __u32 idx;

    memset(stats, 0, sizeof(stats[0]) * 2);

    for (idx = 0; idx < 2; idx++) {
        int ret, cpu;

        ret = bpf_map_lookup_elem(bpf_map__fd(skel->maps.stats),
                      &idx, cnts[idx]);
        if (ret < 0)
            continue;
        for (cpu = 0; cpu < nr_cpus; cpu++)
            stats[idx] += cnts[idx][cpu];
    }
}

/* 发送消息到内核态BPF程序 */
static void send_message_to_kernel(struct scx_hello *skel, const char *message)
{
    u32 key = 0;
    struct scx_hello_message msg;
    int msg_map_fd;
    
    msg_map_fd = bpf_map__fd(skel->maps.hello_message);
    if (msg_map_fd < 0) {
        fprintf(stderr, "无法获取hello_message映射的文件描述符\n");
        return;
    }
    
    /* 先读取当前消息 */
    if (bpf_map_lookup_elem(msg_map_fd, &key, &msg) < 0) {
        fprintf(stderr, "无法读取消息映射\n");
        return;
    }
    
    /* 更新消息 */
    strncpy(msg.message, message, sizeof(msg.message) - 1);
    msg.message[sizeof(msg.message) - 1] = '\0';
    msg.new_message = true;
    
    /* 写回映射 */
    if (bpf_map_update_elem(msg_map_fd, &key, &msg, 0) < 0) {
        fprintf(stderr, "无法更新消息映射\n");
        return;
    }
    
    printf("已发送消息到内核: %s\n", message);
}

int main(int argc, char **argv)
{
    struct scx_hello *skel;
    struct bpf_link *link;
    __u32 opt;
    __u64 ecode;
    bool send_message = false;

    libbpf_set_print(libbpf_print_fn);
    signal(SIGINT, sigint_handler);
    signal(SIGTERM, sigint_handler);

restart:
    skel = SCX_OPS_OPEN(hello_ops, scx_hello);

    while ((opt = getopt(argc, argv, "fvm:i:t:h")) != -1) {
        switch (opt) {
        case 'f':
            skel->rodata->fifo_sched = true;
            break;
        case 'v':
            verbose = true;
            break;
        case 'm':
            strncpy(user_message, optarg, sizeof(user_message) - 1);
            user_message[sizeof(user_message) - 1] = '\0';
            send_message = true;
            break;
        case 'i':
            llc_collect_interval = atoi(optarg);
            if (llc_collect_interval < 1) llc_collect_interval = 1;
            break;
        case 't':
            llc_collect_times = atoi(optarg);
            if (llc_collect_times < 0) llc_collect_times = 0;
            break;
        default:
            fprintf(stderr, help_fmt, basename(argv[0]));
            return opt != 'h';
        }
    }

    SCX_OPS_LOAD(skel, hello_ops, scx_hello, uei);
    link = SCX_OPS_ATTACH(skel, hello_ops, scx_hello);
    
    /* 保存全局skel指针，供LLC采集线程使用 */
    g_skel = skel;
    
    /* 如果指定了消息，发送到内核 */
    if (send_message) {
        send_message_to_kernel(skel, user_message);
    }
    
    /* 创建LLC缓存未命中率采集线程 */
    if (pthread_create(&llc_thread, NULL, llc_collect_thread, NULL) != 0) {
        fprintf(stderr, "创建LLC缓存未命中率采集线程失败\n");
    }

    while (!exit_req && !UEI_EXITED(skel, uei)) {
        __u64 stats[2];
        char new_message[256];

        read_stats(skel, stats);
        printf("local=%llu global=%llu\n", stats[0], stats[1]);
        
        /* 允许用户输入新消息 */
        printf("输入新消息 (直接回车跳过): ");
        fflush(stdout);
        
        if (fgets(new_message, sizeof(new_message), stdin)) {
            /* 移除换行符 */
            size_t len = strlen(new_message);
            if (len > 0 && new_message[len-1] == '\n') {
                new_message[len-1] = '\0';
            }
            
            /* 如果输入不为空，发送消息 */
            if (strlen(new_message) > 0) {
                send_message_to_kernel(skel, new_message);
            }
        }
        
        sleep(1);
    }
    
    /* 等待LLC采集线程结束 */
    pthread_join(llc_thread, NULL);

    bpf_link__destroy(link);
    ecode = UEI_REPORT(skel, uei);
    scx_hello__destroy(skel);

    if (UEI_ECODE_RESTART(ecode))
        goto restart;
    return 0;
}